package main

import (
	"encoding/binary"
	"flag"
	"fmt"
	"log"
	"net"
	"time"

	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
	"github.com/google/gopacket/pcap"
)

const (
	IPPROTO_OSPF   = 89
	OSPFVersion    = 2
	OSPFHelloType  = 1
	OSPFHeaderLen  = 24
	HelloPacketLen = 20
	snapLen        = 65536
	promiscuous    = false
	timeout        = pcap.BlockForever
)

type OSPFClient struct {
	handle    *pcap.Handle
	sourceIP  net.IP
	destIP    net.IP
	routerID  uint32
	ifaceName string
}

func NewOSPFClient(sourceIP, destIP string, ifaceName string) (*OSPFClient, error) {
	// 打开网络接口
	handle, err := pcap.OpenLive(ifaceName, snapLen, promiscuous, timeout)
	if err != nil {
		return nil, fmt.Errorf("打开网络接口失败: %v", err)
	}

	srcIP := net.ParseIP(sourceIP)
	if srcIP == nil {
		handle.Close()
		return nil, fmt.Errorf("无效的源IP地址")
	}

	dstIP := net.ParseIP(destIP)
	if dstIP == nil {
		handle.Close()
		return nil, fmt.Errorf("无效的目标IP地址")
	}

	return &OSPFClient{
		handle:    handle,
		sourceIP:  srcIP.To4(),
		destIP:    dstIP.To4(),
		routerID:  binary.BigEndian.Uint32(srcIP.To4()),
		ifaceName: ifaceName,
	}, nil
}

// 计算校验和
func calculateChecksum(data []byte) uint16 {
	var sum uint32
	for i := 0; i < len(data); i += 2 {
		sum += uint32(data[i]) << 8
		if i+1 < len(data) {
			sum += uint32(data[i+1])
		}
	}
	sum = (sum >> 16) + (sum & 0xffff)
	sum = sum + (sum >> 16)
	return ^uint16(sum)
}

func (c *OSPFClient) createOSPFHello(data string) []byte {
	// OSPF头部(24字节) + Hello包(20字节)
	totalLen := OSPFHeaderLen + HelloPacketLen
	packet := make([]byte, totalLen)

	// OSPF Header
	packet[0] = OSPFVersion                                   // 版本
	packet[1] = OSPFHelloType                                 // 类型
	binary.BigEndian.PutUint16(packet[2:4], uint16(totalLen)) // 长度
	binary.BigEndian.PutUint32(packet[4:8], c.routerID)       // Router ID
	binary.BigEndian.PutUint32(packet[8:12], 0)               // Area ID
	// 校验和位置：packet[12:14]
	binary.BigEndian.PutUint16(packet[14:16], 0) // AuType
	binary.BigEndian.PutUint64(packet[16:24], 0) // Authentication

	// Hello Packet
	binary.BigEndian.PutUint32(packet[24:28], 0xffffff00) // Network Mask
	binary.BigEndian.PutUint16(packet[28:30], 10)         // Hello Interval
	packet[30] = 0x02                                     // Options
	packet[31] = 1                                        // Router Priority
	binary.BigEndian.PutUint32(packet[32:36], 40)         // Dead Interval

	// 在这里嵌入数据
	dataBytes := []byte(data)
	if len(dataBytes) > 4 {
		dataBytes = dataBytes[:4] // 限制数据长度为4字节
	}
	copy(packet[36:40], dataBytes)

	// DR和BDR
	binary.BigEndian.PutUint32(packet[40:44], 0) // Designated Router

	// 计算校验和
	checksum := calculateChecksum(packet)
	binary.BigEndian.PutUint16(packet[12:14], checksum)

	return packet
}

func (c *OSPFClient) SendData(data string) error {
	// 创建以太网层
	eth := layers.Ethernet{
		SrcMAC:       c.getInterfaceMAC(),
		DstMAC:       net.HardwareAddr{0x01, 0x00, 0x5e, 0x00, 0x00, 0x05}, // OSPF组播MAC地址
		EthernetType: layers.EthernetTypeIPv4,
	}

	// 创建IP层
	ip := layers.IPv4{
		Version:  4,
		IHL:      5,
		TOS:      0xc0, // 控制优先级
		TTL:      1,
		Protocol: layers.IPProtocolOSPF,
		SrcIP:    c.sourceIP,
		DstIP:    c.destIP,
	}

	// 创建OSPF数据
	ospfData := c.createOSPFHello(data)

	// 序列化数据包
	buffer := gopacket.NewSerializeBuffer()
	opts := gopacket.SerializeOptions{
		FixLengths:       true,
		ComputeChecksums: true,
	}

	err := gopacket.SerializeLayers(buffer, opts,
		&eth,
		&ip,
		gopacket.Payload(ospfData),
	)
	if err != nil {
		return fmt.Errorf("序列化数据包失败: %v", err)
	}

	// 发送数据包
	if err := c.handle.WritePacketData(buffer.Bytes()); err != nil {
		return fmt.Errorf("发送数据包失败: %v", err)
	}

	return nil
}

func (c *OSPFClient) getInterfaceMAC() net.HardwareAddr {
	iface, err := net.InterfaceByName(c.ifaceName)
	if err != nil {
		return nil
	}
	return iface.HardwareAddr
}

func (c *OSPFClient) Close() {
	if c.handle != nil {
		c.handle.Close()
	}
}

// 列出所有可用的网络接口
func listInterfaces() {
	devices, err := pcap.FindAllDevs()
	if err != nil {
		log.Printf("查找网络接口失败: %v", err)
		return
	}

	fmt.Println("可用的网络接口:")
	for _, device := range devices {
		fmt.Printf("名称: %s\n", device.Name)
		fmt.Printf("描述: %s\n", device.Description)
		fmt.Printf("IP地址: \n")
		for _, address := range device.Addresses {
			fmt.Printf("  %s\n", address.IP)
		}
		fmt.Println()
	}
}

func main() {
	var (
		sourceIP = flag.String("source", "", "源IP地址")
		destIP   = flag.String("dest", "224.0.0.5", "目标IP地址(默认: OSPF组播地址)")
		iface    = flag.String("i", "", "网络接口名称")
		data     = flag.String("data", "TEST", "要发送的数据")
		interval = flag.Int("interval", 10, "发送间隔(秒)")
		list     = flag.Bool("list", false, "列出所有网络接口")
	)
	flag.Parse()

	if *list {
		listInterfaces()
		return
	}

	if *sourceIP == "" || *iface == "" {
		log.Fatal("请指定源IP地址(-source)和网络接口(-i)")
	}

	client, err := NewOSPFClient(*sourceIP, *destIP, *iface)
	if err != nil {
		log.Fatalf("创建OSPF客户端失败: %v", err)
	}
	defer client.Close()

	log.Printf("OSPF客户端已启动")
	log.Printf("源IP: %s", *sourceIP)
	log.Printf("目标IP: %s", *destIP)
	log.Printf("网络接口: %s", *iface)
	log.Printf("发送间隔: %d秒", *interval)

	ticker := time.NewTicker(time.Duration(*interval) * time.Second)
	defer ticker.Stop()

	for range ticker.C {
		if err := client.SendData(*data); err != nil {
			log.Printf("发送数据失败: %v", err)
			continue
		}
		log.Printf("已发送数据: %s", *data)
	}
}
